import matplotlib.pyplot as plt
import os

import numpy as np

struct_dit = '/home/a409/users/lihaowei/data/CVUSA_DSM'
# angle \in (0,360]
# orientation \in (-180, 180]
angle = 180


def val_ind(orientation, angle, w):
    mid = int(w * (0.5 + orientation / 360))
    left_index = int(mid - angle * w / 720)
    right_index = int(mid + angle * w / 720)
    return left_index, right_index


if __name__ == '__main__':
    path = os.path.join(struct_dit, 'splits/val.csv')
    with open(path, 'r') as file:
        line_number = len(file.readlines())
    print(line_number)
    i = 0
    with open(path, 'r') as file:
        for line in file:
            dir_ = line.split(',')[1]
            pic = plt.imread(os.path.join(struct_dit, dir_))

            # h w c = 224,1232,3
            h, w, c = list(pic.shape)
            pic = np.append(pic, pic, axis=1)
            ori = int(np.random.rand() * 361)
            left, right = val_ind(ori, angle, w)
            pic = pic[:, left: right, :]
            save_path = struct_dit + '/angle_' + str(angle) + '_doing'

            os.makedirs(save_path, exist_ok=True)
            plt.imsave(os.path.join(save_path, dir_.split('/')[-1]), pic)

            loss_writer_path = struct_dit + '/angle_ori/'
            os.makedirs(os.path.dirname(loss_writer_path), exist_ok=True)
            file = open(loss_writer_path + '180_val_angle_ori.txt', 'a')
            file.write(dir_ + ',' + str(ori) + '\n')
            i = i + 1
            print('%d/%d' % (i, line_number))




